from skorch.core import *
from skorch.net import NeuralNet
from skorch.helper import predefined_split
from skorch.callbacks import EpochScoring, LRScheduler, CyclicLR, Checkpoint
from torch.optim import SGD
from torch.nn.functional import binary_cross_entropy_with_logits
print(os.getcwd())
sys.path.append(os.getcwd())
from model import UNet
from dataset import create_jsrt_datasets, create_augmenter, jsrt_stats, denormalize_image
torch.manual_seed(0)

class BCEWithLogitsLoss(nn.BCEWithLogitsLoss):
    def __init__(self):
        super(BCEWithLogitsLoss, self).__init__()

    def forward(self, input, target):
        loss = super(BCEWithLogitsLoss, self).forward(input, target)
        # loss2 = super(BCEWithLogitsLossPadding, self).forward(target_1.squeeze_(dim=1), target[1].squeeze_(dim=1))
        # loss3 = super(BCEWithLogitsLossPadding, self).forward(target_1.squeeze_(dim=1), target[2].squeeze_(dim=1))
        return loss



target_size=(512,512)
data_root = '/home/blake/data/dataset/datasets/DR/dicom/DR'
train_ds, valida_ds =create_jsrt_datasets(data_root, create_augmenter, jsrt_stats, target_size=target_size)

module = UNet(n_classes=3, feature_scale=1, imsize=target_size)

# module = model_cuda(module)

net = NeuralNet(
    module,
    criterion=BCEWithLogitsLoss,
    batch_size=4,
    max_epochs=1000,
    optimizer=SGD,
    optimizer__momentum=0.9,
    iterator_train__shuffle=True,
    iterator_train__num_workers=8,
    iterator_valid__shuffle=False,
    iterator_valid__num_workers=8,
    device='cuda',
)
net.initialize()

net.load_params(f_params='best_images/best_params.pt')

val_masks = net.predict(valida_ds)
val_prob_masks = 1/(1 + np.exp(-val_masks))


# We plot the predicted mask with its corresponding true mask and original image:

# In[20]:


def plot_mask_cells(mask_cells):
    """Plots cells with their true mask, predicted mask.

    Parameters
    ----------
    mask_cells: list of tuples (`true_mask`, `predicted_mask`, `cell`)
    padding: int (default=16)
        Padding around mask to remove.
    """
    fig, axes = plt.subplots(len(mask_cells), 7, figsize=(12, 10))
    for idx, (axes, mask_cell) in enumerate(zip(axes, mask_cells), 1):
        ax1, ax2, ax3, ax4, ax5, ax6, ax7 = axes
        true_mask, predicted_mask, cell = mask_cell
        plot_mask_cell(
                true_mask, predicted_mask, cell,
                'Type {}'.format(idx),
                ax1, ax2, ax3, ax4, ax5, ax6, ax7)
    fig.tight_layout()
    return fig, axes


def plot_mask_cell(true_mask,
                   predicted_mask,
                   cell,
                   suffix,
                   ax1, ax2, ax3, ax4, ax5, ax6, ax7):
    """Plots a single cell with a its true mask and predicuted mask"""
    for ax in [ax1, ax2, ax3, ax4, ax5, ax6, ax7]:
        ax.grid(False)
        ax.set_xticks([])
        ax.set_yticks([])

    show_image = denormalize_image(cell, jsrt_stats[0], jsrt_stats[1])
    show_image = np.transpose(show_image, [1,2,0])
    true_mask = to_detach(true_mask, cpu=True)
    true_mask = to_np(true_mask)

    ax1.imshow(show_image)
    ax1.set_title('Image - {}'.format(suffix))

    ax2.imshow(np.squeeze(true_mask[0]), cmap='viridis')
    ax2.set_title('True Mask Lung - {}'.format(suffix))
    ax3.imshow(np.squeeze(true_mask[1]), cmap='viridis')
    ax3.set_title('True Mask Clavicle - {}'.format(suffix))
    ax4.imshow(np.squeeze(true_mask[2]), cmap='viridis')
    ax4.set_title('True Mask Heart - {}'.format(suffix))

    ax5.imshow(predicted_mask[0], cmap='viridis')
    ax5.set_title('Predicted Lung - {}'.format(suffix))
    ax6.imshow(predicted_mask[1], cmap='viridis')
    ax6.set_title('Predicted Clavicle - {}'.format(suffix))
    ax7.imshow(predicted_mask[2], cmap='viridis')
    ax7.set_title('Predicted Heart - {}'.format(suffix))
    return ax1, ax2, ax3, ax4, ax5, ax6, ax7

mask_cells = []
for case_id in [0, 8, 18]:
    cell, mask = valida_ds[case_id]
    mask_cells.append((mask, val_prob_masks[case_id], cell))

_ = plot_mask_cells(mask_cells)

plt.show()